/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: mc2 profiling
 * Author: huawei
 * Create: 2024-05-24
 */

#include <sstream>
#include <unordered_map>
#include "common/sqe_context.h"
#include "task_profiling_pub.h"
#include "common/aicpu_hccl_common.h"
#include "debug/dfx/profiling/profiling_manager.h"
#include "profiling_extend_info.h"

namespace dfx {
uint64_t g_taskHashIds[static_cast<uint64_t>(hccl::ProfTaskType::TASK_INVALID)];
uint64_t g_groupNameHashId{0U};
const std::vector<hccl::ProfTaskType> kfcTaskTypes = {hccl::ProfTaskType::TASK_HCCL_INFO,  // 当前未支持的用这个来暂替
    hccl::ProfTaskType::TASK_NOTIFY_RECORD,
    hccl::ProfTaskType::TASK_NOTIFY_WAIT,
    hccl::ProfTaskType::TASK_SDMA,
    hccl::ProfTaskType::TASK_INTER_RANK_RECORD,
    hccl::ProfTaskType::TASK_INTER_PROCESSOR_SYNC,
    hccl::ProfTaskType::TASK_REDUCE_INLINE,
    hccl::ProfTaskType::TASK_RDMA};

namespace {
void ParseNotifySqeInfo(const SqeInfo &sqeInfo, MsprofAicpuHcclTaskInfo &msprofAicpuMC2HcclInfo)
{
    msprofAicpuMC2HcclInfo.notifyID = sqeInfo.notifyId;
    msprofAicpuMC2HcclInfo.remoteRank = sqeInfo.remoteRank;
    msprofAicpuMC2HcclInfo.itemId = g_taskHashIds[static_cast<uint64_t>(hccl::ProfTaskType::TASK_NOTIFY_RECORD)];
};

void ParseWaitqeInfo(const SqeInfo &sqeInfo, MsprofAicpuHcclTaskInfo &msprofAicpuMC2HcclInfo)
{
    msprofAicpuMC2HcclInfo.notifyID = sqeInfo.notifyId;
    msprofAicpuMC2HcclInfo.remoteRank = sqeInfo.remoteRank;
    msprofAicpuMC2HcclInfo.role = static_cast<uint32_t>(hccl::TaskRole::DST);
    msprofAicpuMC2HcclInfo.itemId = g_taskHashIds[static_cast<uint64_t>(hccl::ProfTaskType::TASK_NOTIFY_WAIT)];
};

void ParseSdmaSqeInfo(const SqeInfo &sqeInfo, MsprofAicpuHcclTaskInfo &msprofAicpuMC2HcclInfo)
{
    // addrXHigh存的是地址的高32位
    msprofAicpuMC2HcclInfo.srcAddr = (static_cast<uint64_t>(sqeInfo.addr1High) << 32) | sqeInfo.addr1Low;
    // addrXLow存的是地址的低32位
    msprofAicpuMC2HcclInfo.dstAddr = (static_cast<uint64_t>(sqeInfo.addr2High) << 32) | sqeInfo.addr2Low;
    msprofAicpuMC2HcclInfo.remoteRank = sqeInfo.remoteRank;
    msprofAicpuMC2HcclInfo.role = static_cast<uint32_t>(hccl::TaskRole::DST);
    if (msprofAicpuMC2HcclInfo.localRank == msprofAicpuMC2HcclInfo.remoteRank) {
        msprofAicpuMC2HcclInfo.transportType = static_cast<uint32_t>(hccl::SimpleTaskType::LOCAL);
    } else {
        msprofAicpuMC2HcclInfo.transportType = static_cast<uint32_t>(hccl::SimpleTaskType::SDMA);
    }
    msprofAicpuMC2HcclInfo.linkType = sqeInfo.taskRelated.linkType;
    msprofAicpuMC2HcclInfo.dataSize = sqeInfo.length;
    if (sqeInfo.opCode == 0) { // 0表示不做随路规约
        // SDMA不展示数据类型
        msprofAicpuMC2HcclInfo.itemId = g_taskHashIds[static_cast<uint64_t>(hccl::ProfTaskType::TASK_SDMA)];
    } else {
        TranslateOpcode(sqeInfo.opCode, msprofAicpuMC2HcclInfo.opType);
        msprofAicpuMC2HcclInfo.itemId = g_taskHashIds[static_cast<uint64_t>(hccl::ProfTaskType::TASK_REDUCE_INLINE)];
    }
}

void ParseCommonSqeInfo(const SqeInfo &sqeInfo, MsprofAicpuHcclTaskInfo &msprofAicpuMC2HcclInfo)
{
    HCCL_WARNING("Unsupported SQE type");
    msprofAicpuMC2HcclInfo.itemId = g_taskHashIds[static_cast<uint64_t>(hccl::ProfTaskType::TASK_HCCL_INFO)];
};

// write value用于卡间的record
void ParseWriteValueSqeInfo(const SqeInfo &sqeInfo, MsprofAicpuHcclTaskInfo &msprofAicpuMC2HcclInfo)
{
    msprofAicpuMC2HcclInfo.remoteRank = sqeInfo.remoteRank;
    if (sqeInfo.subType == RT_STARS_WRITE_VALUE_SUB_TYPE_RDMA_DB_SEND) {
        msprofAicpuMC2HcclInfo.itemId = g_taskHashIds[static_cast<uint64_t>(hccl::ProfTaskType::TASK_RDMA)];
        msprofAicpuMC2HcclInfo.linkType = static_cast<uint32_t>(hccl::LinkType::LINK_ROCE);      // reserved value
        msprofAicpuMC2HcclInfo.transportType = static_cast<uint32_t>(hccl::SimpleTaskType::RDMA); // reserved value
        msprofAicpuMC2HcclInfo.dataSize = sqeInfo.length; // wr len
        msprofAicpuMC2HcclInfo.rdmaType = sqeInfo.taskRelated.rdmaType;
    } else {
        msprofAicpuMC2HcclInfo.itemId = g_taskHashIds[static_cast<uint64_t>(hccl::ProfTaskType::TASK_INTER_RANK_RECORD)];
    }
};

void ParseCondSqeInfo(const SqeInfo &sqeInfo, MsprofAicpuHcclTaskInfo &msprofAicpuMC2HcclInfo)
{
    msprofAicpuMC2HcclInfo.itemId = g_taskHashIds[static_cast<uint64_t>(hccl::ProfTaskType::TASK_INTER_PROCESSOR_SYNC)];
};
}  // namespace

void ProfilingExtendInfoHelper::SqeInfo2MsprofAicpuMC2HcclInfo(
    const SqeInfo &sqeInfo, MsprofAicpuHcclTaskInfo &msprofAicpuMC2HcclInfo)
{
    msprofAicpuMC2HcclInfo.taskId = sqeInfo.taskId;
    msprofAicpuMC2HcclInfo.streamId = sqeInfo.streamId;
    return AssembleProfInfoByType(sqeInfo, msprofAicpuMC2HcclInfo);
}

void ProfilingExtendInfoHelper::AssembleProfInfoByType(
    const SqeInfo &sqeInfo, MsprofAicpuHcclTaskInfo &msprofAicpuMC2HcclInfo)
{
    static const std::unordered_map<uint8_t, Handle> funcMap = {{RT_STARS_SQE_TYPE_WRITE_VALUE, ParseWriteValueSqeInfo},
        {RT_STARS_SQE_TYPE_NOTIFY_RECORD, ParseNotifySqeInfo},
        {RT_STARS_SQE_TYPE_NOTIFY_WAIT, ParseWaitqeInfo},
        {RT_STARS_SQE_TYPE_SDMA, ParseSdmaSqeInfo},
        {RT_STARS_SQE_TYPE_COND, ParseCondSqeInfo}};
    auto it = funcMap.find(sqeInfo.type);
    if (it == funcMap.cend()) {
        return ParseCommonSqeInfo(sqeInfo, msprofAicpuMC2HcclInfo);
    }
    (it->second)(sqeInfo, msprofAicpuMC2HcclInfo);
}
void ProfilingExtendInfoHelper::Ctx2MsprofAicpuMC2HcclInfo(
    const AicpuComContext *ctx, MsprofAicpuHcclTaskInfo &msprofAicpuMC2HcclInfo)
{
    InitHcclInfo(msprofAicpuMC2HcclInfo);
    msprofAicpuMC2HcclInfo.groupName = g_groupNameHashId;
    msprofAicpuMC2HcclInfo.localRank = ctx->rankId;
    msprofAicpuMC2HcclInfo.remoteRank = ctx->rankId;
    msprofAicpuMC2HcclInfo.rankSize = ctx->rankNum;
    msprofAicpuMC2HcclInfo.opType = static_cast<uint32_t>(ctx->reducekind);
}

void ProfilingExtendInfoHelper::InitHcclInfo(MsprofAicpuHcclTaskInfo &msprofAicpuMC2HcclInfo)
{
    msprofAicpuMC2HcclInfo.linkType = static_cast<uint8_t>(hccl::LinkType::LINK_RESERVED);
    msprofAicpuMC2HcclInfo.rdmaType = static_cast<uint8_t>(hccl::RdmaType::RDMA_TYPE_RESERVED);
    msprofAicpuMC2HcclInfo.dataType = static_cast<uint8_t>(HcclDataType::HCCL_DATA_TYPE_RESERVED);
    msprofAicpuMC2HcclInfo.opType = static_cast<uint8_t>(HcclReduceOp::HCCL_REDUCE_RESERVED);
    msprofAicpuMC2HcclInfo.workFlowMode = static_cast<uint8_t>(HcclWorkflowMode::HCCL_WORKFLOW_MODE_RESERVED);
    msprofAicpuMC2HcclInfo.stage = 0;
    msprofAicpuMC2HcclInfo.role = static_cast<uint8_t>(hccl::TaskRole::SRC);
    msprofAicpuMC2HcclInfo.transportType = static_cast<uint8_t>(hccl::SimpleTaskType::LOCAL);
}

std::string ProfilingExtendInfoHelper::AicpuKfcProfCommTurnToString(const AicpuKfcProfCommTurn &aicpuKfcProfCommTurn)
{
    std::stringstream ss;
    ss << "serverStartTime: " << aicpuKfcProfCommTurn.serverStartTime;
    ss << " waitMsgStartTime: " << aicpuKfcProfCommTurn.waitMsgStartTime;
    ss << " kfcAlgExeStartTime: " << aicpuKfcProfCommTurn.kfcAlgExeStartTime;
    ss << " sendTaskStartTime: " << aicpuKfcProfCommTurn.sendTaskStartTime;
    ss << " sendSqeFinishTime: " << aicpuKfcProfCommTurn.sendSqeFinishTime;
    ss << " rtsqExeEndTime: " << aicpuKfcProfCommTurn.rtsqExeEndTime;
    ss << " serverEndTime: " << aicpuKfcProfCommTurn.serverEndTime;
    ss << " dataLen: " << aicpuKfcProfCommTurn.dataLen;
    ss << " deviceId: " << aicpuKfcProfCommTurn.deviceId;
    ss << " streamId: " << aicpuKfcProfCommTurn.streamId;
    ss << " taskId: " << aicpuKfcProfCommTurn.taskId;
    ss << " version: " << aicpuKfcProfCommTurn.version;
    ss << " total commTurn: " << aicpuKfcProfCommTurn.commTurn;
    ss << " currentTurn: " << aicpuKfcProfCommTurn.currentTurn;
    return ss.str();
}

void ProfilingExtendInfoHelper::InitProfItemId()
{
    for (const auto taskType : kfcTaskTypes) {
        // index保证是有效的
        g_taskHashIds[static_cast<uint64_t>(taskType)] =
            AdprofGetHashId(hccl::GetProfTaskOpName(taskType).c_str(), hccl::GetProfTaskOpName(taskType).length());
    }
    return;
}

void ProfilingExtendInfoHelper::Init(const AicpuComContext *ctx)
{
    if (AdprofGetHashId == nullptr) {
        HCCL_INFO("AdprofGetHashId is null, just return");
        return;
    }
    InitProfItemId();
    // ctx保证了非空
    u32 len = 0;                        // 初始化长度为0
    while (ctx->hcomId[len] != '\0') {  // 遍历数组，直到遇到'\0'
        len++;                          // 每遍历一个字符，长度加1
    }
    g_groupNameHashId = AdprofGetHashId(ctx->hcomId, len);
    HCCL_DEBUG("Using %s and len %u to get hash %lu", ctx->hcomId, len, g_groupNameHashId);
}
}  // namespace dfx